import random
import math

def assign_cluster(x, centroids):
    min_dist = float('inf')  
    cluster_idx = 0         
    
    for idx, centroid in enumerate(centroids):
        # 计算数据点x与质心centroid的欧氏距离
        dist = 0.0
        for xi, ci in zip(x, centroid):
            dist += (xi - ci) ** 2  
        dist = math.sqrt(dist)      # 欧氏距离
        if dist < min_dist:
            min_dist = dist
            cluster_idx = idx
    
    return cluster_idx

def Kmeans(data, k, epsilon=1e-3, iteration=100):
    # 输入合法性校验
    if not data or len(data) <= k:
        raise ValueError("数据集长度必须大于聚类数量k")
    if k <= 1:
        raise ValueError("聚类数量k必须大于1")
    dim = len(data[0]) 
    for x in data:
        if len(x) != dim:
            raise ValueError("所有数据点必须具有相同的维度")
    
    min_vals = [min(x[d] for x in data) for d in range(dim)]
    max_vals = [max(x[d] for x in data) for d in range(dim)]
    
    centroids = []
    for _ in range(k):
        centroid = [random.uniform(min_vals[d], max_vals[d]) for d in range(dim)]
        centroids.append(centroid)
    
    iter_count = 0
    while iter_count < iteration:
        old_centroids = [c.copy() for c in centroids]
        
        clusters = []  
        for x in data:
            cluster_idx = assign_cluster(x, centroids)
            clusters.append(cluster_idx)
        
        sum_clusters = [[0.0 for _ in range(dim)] for _ in range(k)]  # 各聚类各维度总和
        count_clusters = [0 for _ in range(k)]                        # 各聚类数据点数量
        
        for x, idx in zip(data, clusters):
            for d in range(dim):
                sum_clusters[idx][d] += x[d]
            count_clusters[idx] += 1
        
        for i in range(k):
            if count_clusters[i] > 0:
                centroids[i] = [sum_clusters[i][d] / count_clusters[i] for d in range(dim)]

        max_centroid_change = 0.0
        for old_c, new_c in zip(old_centroids, centroids):
            change = 0.0
            for oc, nc in zip(old_c, new_c):
                change += (oc - nc) ** 2
            change = math.sqrt(change)
            if change > max_centroid_change:
                max_centroid_change = change

        if max_centroid_change < epsilon:
            print(f"迭代{iter_count+1}次后收敛（质心最大变化量：{max_centroid_change:.6f} < {epsilon}）")
            break
        
        iter_count += 1
    
    if iter_count >= iteration:
        print(f"已达到最大迭代次数{iteration}，未完全收敛（质心最大变化量：{max_centroid_change:.6f}）")
    
    return clusters, centroids

if __name__ == "__main__":
    print("="*50)
    print("测试1：一维数据聚类（模拟灰度值）")
    gray_data = [[12], [15], [18], [20], [85], [90], [92], [95], [98], [100]]  
    clusters1, centroids1 = Kmeans(data=gray_data, k=2, epsilon=1e-4, iteration=50)
    print(f"原始数据：{[x[0] for x in gray_data]}")
    print(f"聚类结果（0/1表示聚类索引）：{clusters1}")
    print(f"最终质心：{[round(c[0], 2) for c in centroids1]}")  
    print("\n" + "="*50)
    print("测试2：二维数据聚类")
    two_d_data = [
        [1.0, 2.0], [1.5, 1.8], [5.0, 8.0], [8.0, 8.0],
        [1.0, 0.6], [9.0, 11.0], [8.0, 2.0], [10.0, 2.0],
        [9.0, 3.0], [0.5, 1.0], [7.0, 9.0], [6.0, 8.5]
    ] 
    clusters2, centroids2 = Kmeans(data=two_d_data, k=3, epsilon=1e-4, iteration=50)
    print(f"二维数据聚类结果：{clusters2}")
    print(f"最终质心：{[(round(c[0],2), round(c[1],2)) for c in centroids2]}")
